from Logging import Logs, Log
from Visualizer  import Visualizer
from Analyzer import Analyzer
from utils import Operation, debug
import math

class backtest():
    def __init__(self,
                 data,
                 Strategy,
                 verbose=True,
                 logname=None,
                 start_day='2010-01-04',
                 end_day = '2019-12-31',
                 visualization=True):
        self.data = data
        self.Strategy = Strategy
        self.start_day = start_day
        self.end_day = end_day
        self.Logs = Logs(logname, Strategy.name)
        self.Visualizer = Visualizer(self.Logs.logs) if visualization else None
        self.Analyzer = Analyzer(self.Logs.logs)
        self.verbose = verbose

        self.dates = list(self.data.keys())
        self.dates.sort()
        self.ids = [i for i in self.data[self.dates[0]].index]
        self.items = [i for i in data[self.dates[0]].columns]
        assert start_day in self.dates and end_day in self.dates
        assert self.Strategy.every_n_day > 0


    def overall_value(self, date, holdings, money):
        sum = money
        for k in holdings.keys():
            sum += self.stock_price(k, date) * holdings[k]
        return sum

    def stock_price(self, id, date):
        return self.data[date]['close'][id]

    def operate_price(self, id, date):
        return self.data[date]['open'][id]

    def execute_op(self, op, date):
        id, value, ratio = op.id, op.value, op.ratio
        if op.op == 'none':
            pass
        elif op.op == 'sell':
            absolute = op.value * self.Strategy.holdings[id] if op.ratio else op.value
            if absolute > self.Strategy.holdings[id]:
                Exception("cant sell that much")
            self.Strategy.money += absolute * self.operate_price(id, date)
            if self.Strategy.holdings[id] == absolute:
                self.Strategy.holdings.pop(id)
            else:
                self.Strategy.holdings[id] -= absolute
            if self.verbose:
                print("sold", op.id, absolute, "x", "price:", self.operate_price(id, date),'got:',round(absolute * self.operate_price(id, date)), 2)
        elif op.op == 'buy':
            cash = self.Strategy.money * op.value if op.ratio else op.value
            stocks = cash / self.operate_price(op.id, date)
            if id in self.Strategy.holdings.keys():
                self.Strategy.holdings[id] += stocks
            else:
                self.Strategy.holdings[id] = stocks
            self.Strategy.money -= cash
            if self.verbose:
                print('bought', op.id, stocks, "x","price:",self.operate_price(id, date), "cost:", round(cash, 2))
        else:
            Exception("operation has wrong op field:", op.op)

    def buy_max(self, id, money, date):
        stock_price = self.operate_price(id, date)
        holds = money // stock_price
        value = holds * stock_price
        return holds, value

    def check_valid(self, id, date):
        if math.isnan(self.stock_price(id, date)):
            return False
        else:
            return True


    def modify_execute_ops(self, ops, date):
        modified_ops = []
        for op in ops:
            if not isinstance(op, Operation):
                debug("strategy didnt return Operation list")
            else:
                if op.op == 'none':
                    modified_ops.append(Operation('none'))
                elif not self.check_valid(op.id, date):
                    continue
                elif op.op == 'sell':
                    if op.id in self.Strategy.holdings.keys():
                        if (not op.ratio and op.value > self.Strategy.holdings[op.id]) or (op.ratio and op.value > 1):
                            modified = Operation('sell', op.id, 1, True)
                        else:
                            modified = op
                        self.execute_op(modified, date)
                        modified_ops.append(modified)
                elif op.op == 'buy':
                    val = self.Strategy.money * op.value if op.ratio else op.value
                    if val > self.Strategy.money:
                        val = self.Strategy.money
                    holds, value = self.buy_max(op.id, val, date)
                    if holds == 0:
                        continue
                    else:
                        modified = Operation('buy', op.id, value, False)
                        self.execute_op(modified, date)
                        modified_ops.append(op)
                else:
                    Warning("strategy returned wrong operation which is not in [sell, buy, none], will be ignored")
        return modified_ops

    def observe(self, ob, date):
        date_idx = self.dates.index(date)
        id, n_days_before, item = ob.id, ob.n_days_before, ob.item
        if date_idx < n_days_before:
            return None
        if item not in self.items:
            Exception("observation request has wrong item:", item, "which is not in", self.items)
        if n_days_before < 1:
            Exception("observation request wants to see the future")
        if id not in self.ids:
            Exception("observation request has wrong id")
        return self.data[self.dates[date_idx-n_days_before]][item][id]

    def fill_observe(self, obs, date):
        self.Strategy.Observed = [self.observe(ob, date) for ob in obs]

    def run(self):
        start_day_idx = self.dates.index(self.start_day)
        end_day_idx = self.dates.index(self.end_day)

        d = start_day_idx
        while(d < end_day_idx):
            self.fill_observe(self.Strategy.Observations, self.dates[d])
            self.Strategy.strategy()
            ops = self.Strategy.Operations
            real_ops = self.modify_execute_ops(ops, self.dates[d])
            self.Logs.record(Log(self.dates[d],
                                 self.Strategy.money,
                                 self.Strategy.holdings,
                                 self.overall_value(self.dates[d], self.Strategy.holdings, self.Strategy.money),
                                 real_ops))
            d += self.Strategy.every_n_day

        print(self.Strategy.name, "backtesting done, logging...")
        self.Logs.output()

    def visualize(self):
        self.Visualizer.visualization()

    def analyze(self):
        self.Analyzer.analysis()

